clear; clc; close all;

I = imread('Lena.tif');
if size(I,3) == 3
    I = rgb2gray(I);
end
I = imresize(I, [256 256]);
I = im2double(I);

% 返回加密图和密钥
[I_encrypted, key] = encrypt_image(I);

imwrite(I_encrypted, 'Lena_encrypted.tif');

% 解密时用同一个key
I_decrypted = decrypt_image(I_encrypted, key);

figure;
subplot(1,3,1); imshow(I); title('Original Image');
subplot(1,3,2); imshow(I_encrypted); title('Encrypted Image');
subplot(1,3,3); imshow(I_decrypted); title('Decrypted Image');

%% ———————— 加密函数及子函数 ————————
function [I_encrypted, key] = encrypt_image(I)
    % 输入：I - 原始图像（任意大小灰度或彩色）
    % 输出：
    %   I_encrypted - 加密后的uint8图像矩阵
    %   key - 结构体，包含解密需要的所有密钥信息

    % --- 预处理：确保256x256灰度double，范围[0,1] ---
    if size(I,3) == 3
        I = rgb2gray(I);
    end
    I = imresize(I, [256 256]);  % 缩放图像到256x256，方便算法统一处理
    I = im2double(I);            % 转double并归一化到[0,1]，保持数值稳定

    % 1. DWT低频子带
    [LL,~,~,~] = dwt2(I,'haar');
    I1_vec = LL(:);

    % 2. 压缩采样参数
    compression_ratio = 0.8;
    vec_len = numel(I1_vec);
    M = round(compression_ratio * vec_len);
    M = floor(sqrt(M))^2;    % 保证M为完全平方数
    side_len = sqrt(M);

    % 3. 超混沌系统参数
    params.a = 1; params.b = 2*pi; params.c = 11.5;
    params.x0 = 0.4; params.y0 = 0.8; params.z0 = 0.6;

    % 4. 生成超混沌序列（多生成200步丢弃）
    [Xall, Yall, Zall] = generate_3D_SIMM_sequences(M+200, params);
    X = Xall(201:end); Y = Yall(201:end); Z = Zall(201:end);

    % 5. 构造Hadamard测量矩阵
    N_hadamard = 2^nextpow2(vec_len);
    H = hadamard(N_hadamard);
    H = H(:,1:vec_len);
    [~, idx] = sort(X, 'descend');
    Phi = H(idx(1:M), :);

    % 6. 压缩采样 + 量化
    I2_vec = Phi * I1_vec;
    I2q = uint8(255 * mat2gray(I2_vec));  % mat2gray保证映射到[0,1]

    % 7. Arnold置乱参数
    a_scramble = mod(floor(abs(X*1e6)), 256);
    b_scramble = mod(floor(abs(Y*1e6)), 256);

    % 8. 置乱
    scrambled_vec = improved_arnold_scramble(I2q, a_scramble(1), b_scramble(1));

    % 9. 转换为矩阵并转置
    scrambled_mat = reshape(scrambled_vec, side_len, side_len)';
    scrambled_vec = scrambled_mat(:);

    % 10. 像素扩散序列
    S1 = mod(floor(abs(Y*1e6)),256);
    S2 = mod(floor(abs(Z*1e6)),256);

    % 11. 像素扩散
    encrypted_vec = diffusion(scrambled_vec, S1, S2);

    % 12. 输出加密图像矩阵
    I_encrypted = reshape(encrypted_vec, side_len, side_len);

    % 13. 返回密钥结构体
    key.Phi = Phi;
    key.S1 = S1(1:length(encrypted_vec));
    key.S2 = S2(1:length(encrypted_vec));
    key.a_scramble = a_scramble(1);
    key.b_scramble = b_scramble(1);
    key.side_len = side_len;
    key.params = params;
end
 

function [X,Y,Z] = generate_3D_SIMM_sequences(M, p)
    a = p.a; b = p.b; c = p.c;
    x0 = p.x0; y0 = p.y0; z0 = p.z0;
    X = zeros(1,M); Y = zeros(1,M); Z = zeros(1,M);
    X(1) = x0; Y(1) = y0; Z(1) = z0;
    for i=1:M-1
        X(i+1) = a*sin(b*Z(i))*sin(c/X(i));
        Y(i+1) = a*sin(b*X(i+1))*sin(c/Y(i));
        Z(i+1) = a*sin(b*Y(i+1))*sin(c/Z(i));
    end
end

function scrambled = improved_arnold_scramble(vec, a, b)
    N = numel(vec);
    scrambled = zeros(size(vec));
    for j=1:N
        q = mod(b + (a*b + 1)*j, N);
        idx = mod(q, N) + 1;
        scrambled(idx) = vec(j);
    end
    scrambled = uint8(scrambled);
end

function out_vec = diffusion(in_vec, S1, S2)
    N = length(in_vec);
    C = zeros(size(in_vec));
    C(1) = mod(double(in_vec(1)) + double(S1(1)), 256);
    for i=2:N
        C(i) = mod(C(i-1) + double(in_vec(i)) + double(S1(i)), 256);
    end
    C2 = zeros(size(C));
    C2(end) = C(end);
    for i=N-1:-1:1
        C2(i) = mod(C2(i+1) + double(C(i)) + double(S2(i)), 256);
    end
    out_vec = uint8(C2);
end



function I_decrypted = decrypt_image(I_encrypted, key)
    side_len = key.side_len;
    encrypted_vec = I_encrypted(:);

    % 1. 逆扩散
    diffused_vec = inverse_diffusion(encrypted_vec, key.S1, key.S2);

    % 2. 逆Arnold置乱
    unscrambled_vec = inverse_arnold_scramble(diffused_vec, key.a_scramble, key.b_scramble, side_len, 10); % 10次置乱，和加密时保持一致

    % 3. OMP稀疏重构
    y = double(unscrambled_vec);
    Phi = key.Phi;
    s = 5000;  % 稀疏度，需与加密时匹配
    N2 = size(Phi, 2);

    x_hat = zeros(N2, 1);
    r = y;
    idx_set = [];
    for t = 1:s
        proj = abs(Phi' * r);
        [~, idx] = max(proj);
        idx_set = [idx_set; idx];
        Phi_sub = Phi(:, idx_set);
        x_ls = Phi_sub \ y;
        r = y - Phi_sub * x_ls;
        if norm(r) < 1e-3
            break;
        end
    end
    x_hat(idx_set) = x_ls;

    % 4. 重构系数转矩阵
    N = sqrt(N2);
    I1_rec = reshape(x_hat, N, N);

    % 5. IDWT反变换
    I_decrypted = idwt2(I1_rec, [], [], [], 'haar');
    I_decrypted = mat2gray(I_decrypted);
end

function out_vec = inverse_diffusion(encrypted_vec, S1, S2)
    N = length(encrypted_vec);
    S1 = S1(1:N);
    S2 = S2(1:N);

    C = zeros(size(encrypted_vec));
    C(1) = encrypted_vec(1);
    for i = 2:N
        C(i) = mod(double(encrypted_vec(i)) - double(encrypted_vec(i-1)) - double(S1(i)), 256);
    end

    D = zeros(size(C));
    D(end) = C(end);
    for i = N-1:-1:1
        D(i) = mod(double(C(i)) - double(C(i+1)) - double(S2(i)), 256);
    end

    out_vec = uint8(D);
end

function out_vec = inverse_arnold_scramble(vec, a, b, M, iter)
    N = length(vec);
    temp_vec = vec;

    for k = 1:iter
        temp = zeros(size(temp_vec));
        for idx = 1:N
            [x, y] = ind2sub([M M], idx);
            % 逆Arnold变换矩阵示例（a=1,b=1）：
            x_new = mod(x - y, M) + 1;
            y_new = mod(-x + 2*y, M) + 1;
            new_idx = sub2ind([M M], x_new, y_new);
            temp(new_idx) = temp_vec(idx);
        end
        temp_vec = temp;
    end

    out_vec = temp_vec;
end
